# leader_system_estimation.py

import numpy as np
import matplotlib.pyplot as plt
from scipy.linalg import pinv

# ========================== 全局参数 ==========================
dt = 0.1        # 采样间隔 (s)
T = 60          # 历史轨迹时长 (s)
tau = int(T / dt) + 1  # 数据点数
np.random.seed(42)  # 保证可复现

# ========================== 1. 生成球坐标轨迹 ==========================
def generate_leader_measurements():
    t = np.linspace(0, T, tau)
    r = 1000 + 50 * np.sin(0.1 * t)
    r_dot = 5 * np.cos(0.1 * t)
    elev = np.deg2rad(10 + 5 * np.sin(0.05 * t))
    elev_dot = np.deg2rad(0.5 * np.cos(0.05 * t))
    azim = np.deg2rad(30 + 10 * np.sin(0.03 * t))
    azim_dot = np.deg2rad(0.3 * np.cos(0.03 * t))
    
    return np.vstack([t, np.zeros_like(t), r, elev, azim, r_dot, elev_dot, azim_dot]).T

# ========================== 2. 球坐标转笛卡尔坐标 ==========================
def spherical_to_cartesian(meas):
    _, _, r, elev, azim, r_dot, elev_dot, azim_dot = meas.T
    x = r * np.cos(elev) * np.cos(azim)
    y = r * np.cos(elev) * np.sin(azim)
    z = r * np.sin(elev)
    
    x_dot = r_dot * np.cos(elev) * np.cos(azim) - r * elev_dot * np.sin(elev) * np.cos(azim) - r * azim_dot * np.cos(elev) * np.sin(azim)
    y_dot = r_dot * np.cos(elev) * np.sin(azim) - r * elev_dot * np.sin(elev) * np.sin(azim) + r * azim_dot * np.cos(elev) * np.cos(azim)
    z_dot = r_dot * np.sin(elev) + r * elev_dot * np.cos(elev)
    
    return np.vstack([x, x_dot, y, y_dot, z, z_dot]).T  # shape = (τ, 6)

# ========================== 3. 输出数据生成 ==========================
def generate_input():
    return np.random.randn(tau - 1, 3) * 2  # shape = (τ-1, 3)

# ========================== 4. 数据驱动估计 S1 ==========================
def estimate_S1(leader_state, U):
    γ = leader_state.T               # shape = (6, τ)
    γ_L = γ[:, :-1]                  # shape = (6, τ-1)
    γ_R = γ[:, 1:]                   # shape = (6, τ-1)
    U_L = U.T                        # shape = (3, τ-1)

    M = np.vstack([γ_L, U_L])        # shape = (9, τ-1)
    rank_M = np.linalg.matrix_rank(M)
    n_γ, n_u = γ_L.shape[0], U_L.shape[0]

    print(f"[INFO] Rank([γ;u]) = {rank_M}, expected = {n_γ + n_u}")
    if rank_M < n_γ + n_u:
        print("⚠️ 数据不满足信息性条件，估计结果可能不唯一或不准确")

    S_all = γ_R @ pinv(M)            # shape = (6, 9)
    S1 = S_all[:, :n_γ]              # shape = (6, 6)
    B_i1 = S_all[:, n_γ:]            # shape = (6, 3)

    return S1, B_i1

# ========================== 5. 可视化 ==========================
def visualize(S1):
    eigs = np.linalg.eigvals(S1)
    print("S1 Eigenvalues:", np.round(eigs, 4))
    plt.subplot(141)  # 1行4列的第1个子图
    plt.scatter(np.real(eigs), np.imag(eigs), marker='x', color='r')
    plt.title("Eigenvalues of S1")
    plt.xlabel("Real")
    plt.ylabel("Imag")
    plt.grid(True)
    plt.axhline(0, color='k', ls='--')
    plt.axvline(0, color='k', ls='--')

def plot_3d_trajectory(positions, title="3D Trajectory", subplot_pos=None):
    if subplot_pos:
        ax = plt.subplot(subplot_pos, projection='3d')
    else:
        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')
    ax.plot(positions[:, 0], positions[:, 1], positions[:, 2], 'b')
    ax.set_xlabel("X (m)")
    ax.set_ylabel("Y (m)")
    ax.set_zlabel("Z (m)")
    ax.set_title(title)

# ========================== 6. 状态递推 v(k+1) = S1*v(k) ==========================
def simulate_leader_from_S1(S1, v0, steps):
    n = S1.shape[0]
    v = np.zeros((steps + 1, n))
    v[0] = v0
    for k in range(steps):
        v[k + 1] = S1 @ v[k]
    return v

# ========================== 7. 主函数 ==========================
def main():
    # 数据生成
    measurements = generate_leader_measurements()
    leader_state = spherical_to_cartesian(measurements)
    U = generate_input()

    # 数据驱动估计
    S1, B_i1 = estimate_S1(leader_state, U)

    # 输出
    print("\n=== S1 估计结果 ===")
    print(np.round(S1, 4))
    print("\n=== B_i1 估计结果 ===")
    print(np.round(B_i1, 4))

    # 创建一个更宽的图形窗口
    plt.figure(figsize=(20, 5))
    
    # S1特征值图
    visualize(S1)

    # 历史轨迹图 (0-60s)
    xyz_hist = leader_state[:, [0, 2, 4]]
    plot_3d_trajectory(xyz_hist, title="True Trajectory\n(0–60s)", subplot_pos=142)

    # 历史轨迹预测图 (0-60s)
    v0 = leader_state[0, :6]
    steps_60 = int(60 / dt)
    v_pred_60 = simulate_leader_from_S1(S1, v0, steps_60)
    xyz_pred_60 = v_pred_60[:, [0, 2, 4]]
    plot_3d_trajectory(xyz_pred_60, title="Predicted Trajectory\n(0–60s)", subplot_pos=143)

    # 未来轨迹预测图 (61-120s)
    steps_120 = int(120 / dt)
    v_pred_120 = simulate_leader_from_S1(S1, v0, steps_120)
    xyz_pred_120 = v_pred_120[steps_60+1:, [0, 2, 4]]  # 只取61-120s的数据
    plot_3d_trajectory(xyz_pred_120, title="Future Trajectory\n(61–120s)", subplot_pos=144)

    # 调整布局并显示
    plt.tight_layout()
    plt.show()

 # ======================== 误差分析 ========================
    # 只取前 60s 的预测状态，与真实状态对齐
    v_pred = v_pred_60  # 使用前60s的预测轨迹
    v_pred_short = v_pred[:tau, :]  # shape = (601, 6)
    v_true = leader_state[:, :6]    # shape = (601, 6)

    pos_pred = v_pred_short[:, [0, 2, 4]]
    pos_true = v_true[:, [0, 2, 4]]
    pos_error = pos_pred - pos_true
    error_norm = np.linalg.norm(pos_error, axis=1)  # 每一时刻的位置误差

    print("\n=== 预测误差统计（前60s）===")
    print(f"平均位置误差: {np.mean(error_norm):.4f} m")
    print(f"最大位置误差: {np.max(error_norm):.4f} m")

    # 绘制误差随时间变化曲线
    t = np.linspace(0, T, tau)
    plt.figure()
    plt.plot(t, error_norm, label="||x_pred - x_true||")
    plt.xlabel("Time (s)")
    plt.ylabel("Position Error (m)")
    plt.title("Prediction Error over Time (0–60s)")
    plt.grid(True)
    plt.legend()
    plt.tight_layout()
    plt.show()
if __name__ == "__main__":
    main()
